//
//  maxHeap.swift
//  heap
//
//  Created by lhc on 2024/8/8.
//

import Foundation


class MaxHeap<T: Comparable> {
    // 最大堆内部实现是一个数组
    
    private var size: Int
    private var dataArray: [T]
    
    enum Exception: Error {
        ///   index 不能小于等于0
        case TheIndexMustNotBeLessThanOrEqualToZero
        ///   index 不能小于0
        case TheIndexMustNotBeLessThanZero
        ///   index 要小于size
        case TheIndexMustBeLessThanSize
        ///   当前堆是一个空堆
        case EmptyHeap
    }
    
    init() {
        self.dataArray = [T]()
        self.size = 0
    }
    
    /// 根据传来的数组生成最大堆
    init(_ arr: [T]) throws {
        self.dataArray = arr
        self.size = self.dataArray.count
        // 如果只有一个元素就不用heapify了
        guard self.size > 1 else {return}
        
        // 获取最后一个父节点的index
        var i = (self.size - 2) / 2
        while i >= 0 {
            try siftDown(i)
            i -= 1
        }
    }
    
    func isEmpty() -> Bool {
        return self.size == 0
    }
    
    func getSize() -> Int {
        return self.size
    }
    
    func getItemByIndex(_ index: Int) throws -> T {
        guard index < self.size else {
            throw Exception.TheIndexMustBeLessThanSize
        }
        
        guard index >= 0 else {
            throw Exception.TheIndexMustNotBeLessThanZero
        }
        
        return self.dataArray[index]
    }
    
    /// 返回父节点index
    private func parentIndexFromIndex(_ index: Int) throws -> Int {
        guard index > 0 else {
            throw Exception.TheIndexMustNotBeLessThanOrEqualToZero
        }
        
        return (index - 1) / 2
    }
    
    /// 返回左孩子的index
    private func leftChildIndexFromIndex(_ index: Int) -> Int {
        return index * 2 + 1
    }
    
    /// 返回右孩子的index
    private func rightChildIndexFromIndex(_ index: Int) -> Int {
        return index * 2 + 2
    }
    
    /// 添加元素
    func add(element: T) {
        self.dataArray.append(element)
        
        self.size += 1
        
        self.siftUp(self.size - 1)
        
    }
    
    /// 把最末尾的元素按照大小提升到他应该的位置
    private func siftUp(_ index: Int = 0) {
        // 获取新天际的最末尾的元素
        let e = self.dataArray[index]
        // 末尾元素的index作为初始值
        var i = index
        while true {
            guard let parentIndex = try? self.parentIndexFromIndex(i) else {
                break
            }
            
            guard e > self.dataArray[parentIndex] else {
                break
            }
            
            try? self.exchange(l: i, r: parentIndex)
            i = parentIndex
        }
        
    }
    
    /// 交换两个index的元素
    private func exchange(l: Int, r: Int) throws {
        guard l >= 0, r >= 0  else {
            throw Exception.TheIndexMustNotBeLessThanZero
        }
        
        guard l < self.size, r < self.size else {
            throw Exception.TheIndexMustBeLessThanSize
        }
        
        let temp = self.dataArray[l]
        self.dataArray[l] = self.dataArray[r]
        self.dataArray[r] = temp
    }
    
    /// 返回最大值
    func findMax() throws -> T {
        guard !self.isEmpty() else {
            throw Exception.EmptyHeap
        }
        
        return self.dataArray.first!
    }
    
    /// 提取最大值
    func extractMax() throws -> T {
        let e = try findMax()
        
        /// 先交换最大值和最末尾的值, 把最大值放到最末尾
        try self.exchange(l: 0, r: self.size - 1)
        /// 删除末尾最大值
        let _ = self.dataArray.removeLast()
        self.size -= 1
        
        try self.siftDown()
        
        return e
    }
    
    private func siftDown(_ index: Int = 0) throws {
        
        var i = index
        // 如果左孩子的index还没超出size，说明左孩子还存在，当前index不是叶结点
        while self.leftChildIndexFromIndex(i) < self.size {
            var k = self.leftChildIndexFromIndex(i)
            // lIndex + 1的index就是右孩子,如果不超过size表示存在
            if k + 1 < self.size && self.dataArray[k] < self.dataArray[k + 1] {
                k = k + 1
            }
            // 这时候k一定是左右孩子中最大值的index
            if self.dataArray[k] > self.dataArray[i] {
                try self.exchange(l: i, r: k)
                i = k
                // 进入下一个循环
            }else {
                break
                // 循环结束
            }
        }
        
    }
    
    /// 替换头结点元素，然后重新实现最大堆
    func replace(element: T) throws -> T {
        let ret = try self.findMax()
        
        self.dataArray[0] = element
        try self.siftDown()
        
        return ret
    }
}

extension MaxHeap: CustomStringConvertible {
    var description: String {
        return "\(self.dataArray)"
    }
    
    
}
